﻿#include <iostream>
#include <queue>
#include "Graph.h"

static struct EdgeComparator
{
	bool operator()(const GraphEdge* a, const GraphEdge* b)
	{
		return a->weight > b->weight;
	}
};

static GraphNode* anotherNode(GraphEdge* edge, GraphNode* oneNode)
{
	return edge->from == oneNode ? edge->to : edge->from;
}

#define EXISTS(SET, X) (SET.find(X) != SET.end())

#define CLEAR_HEAP(HEAP) int size = HEAP.size(); for (int i = 0; i < size; i++) { HEAP.pop(); }

// 最小生成树Prim算法
static void prim(const Graph* graph, unordered_set<GraphEdge*>& edgeResults)
{
	unordered_set<GraphNode*> usedNodes;
	unordered_set<GraphEdge*> backupEdges;
	priority_queue<GraphEdge*, vector<GraphEdge*>, EdgeComparator> sortHeap;
	auto node = graph->nodes.begin()->second;
	usedNodes.insert(node);

	while (usedNodes.size() < graph->nodes.size())
	{
		CLEAR_HEAP(sortHeap);

		for (auto edgeIterator = node->edges.begin(); edgeIterator != node->edges.end(); edgeIterator++)
		{
			auto edge = *edgeIterator;
			if (!EXISTS(edgeResults, edge))
			{
				sortHeap.push(edge);
				backupEdges.insert(edge);
			}
		}

		GraphEdge* minWeightEdge = nullptr;
		do 
		{
			minWeightEdge = sortHeap.top();
			sortHeap.pop();
		} while (!sortHeap.empty() && EXISTS(usedNodes, minWeightEdge->from) && EXISTS(usedNodes, minWeightEdge->to));

		if (minWeightEdge && (!EXISTS(usedNodes, minWeightEdge->from) || !EXISTS(usedNodes, minWeightEdge->to)))
		{
			// 发现最小边的另外一个点可以延续查找
			node = anotherNode(minWeightEdge, node);
			usedNodes.insert(node);
			edgeResults.insert(minWeightEdge);
			backupEdges.erase(minWeightEdge);
		}
		else
		{
			// 从备选边中选择一个weight最小，且有一个没有使用过的端点
			CLEAR_HEAP(sortHeap);

			for (auto edgeIterator = backupEdges.begin(); edgeIterator != backupEdges.end(); ++edgeIterator)
			{
				auto edge = *edgeIterator;
				if (EXISTS(usedNodes, edge->from) && EXISTS(usedNodes, edge->to)) continue;

				sortHeap.push(edge);
			}

			minWeightEdge = sortHeap.top();
			node = EXISTS(usedNodes, minWeightEdge->from) ? minWeightEdge->to : minWeightEdge->from;

			usedNodes.insert(node);
			edgeResults.insert(minWeightEdge);
			backupEdges.erase(minWeightEdge);
		}
	}
}

/**
 * 最小生成树，Prim算法
 * 要求：无向图
 *
 *                A
 *             / |  \
 *            6  1    5
 *           /   |     \
 *          B--5-C--5---D
 *          \    /\     /
 *           3  6  4   2
 *            \/    \/
 *             E--6--F
 *
 * 因为是无向图，故 A-B = B-A
 * 1. 准备一个Node队列Queue，用于存放处理过的节点；准备一个Edge的Set，用于存放处理过的边；准备一个Edge的Set2，用于存放已经解锁的路
 * 2. 从A出发，Queue加入A， 然后解锁A-B，A-C，A-D三条路，选weight最小的A-C，Queue加入C，Set加入A-C，Set2加入A-B,A-C,A-D
 * 3. 从C出发，有C-B, C-D, C-E, C-F, C-A五条路，但是C-A已经在Set中了，无法重新使用，所以选除C-A外weight最小的C-F, Queue加入F，
 * Set加入C-F, Set2加入C-B，C-D，C-E，C-F（C-A已经加入过了）
 * 4. 从F出发，有F-C，F-D，F-E三条路，选weight最小的F-D，Queue加入D，Set加入F-D，Set2加入F-D，F-E（F-C已经加入过了）
 * 5. 从D出发，有D-A，D-C，D-F三条路，但是这三条路的另外一个点A,C,F已经在Queue中了
 * 6. 从Set2中找一条边，这条边中包括不在Queue中的点，weight最小，并且不在Set中的边，找到了C-B，Queue加入B, Set加入C-B
 * 7. 从B出发，有B-A，B-E两条路，选weight最小的B-E，Queue加入E，Set加入B-E, Set2加入B-E
 * 8. 至此，Queue中已经包含了所有的点，Set中的边即为所求。
 */
int main_Prim()
{
	int matrix[][3] = {
		{'a', 'b', 6},
		{'a', 'c', 1},
		{'a', 'd', 5},
		{'b', 'c', 5},
		{'b', 'e', 3},
		{'c', 'd', 5},
		{'c', 'e', 6},
		{'c', 'f', 4},
		{'d', 'f', 2},
		{'e', 'f', 6},
	};

	int* pMatrix = matrix[0];
	Graph* graph = Graph::generate(&pMatrix, sizeof(matrix) / sizeof(int) / 3, 3, false);

	unordered_set<GraphEdge*> edgeResults;
	prim(graph, edgeResults);

	for (auto edgeIterator = edgeResults.begin(); edgeIterator != edgeResults.end(); ++edgeIterator)
	{
		auto edge = *edgeIterator;
		printf("%c->%c %d\n", edge->from->value, edge->to->value, edge->weight);
	}

	delete(graph);
	return 0;
}